import cv2
import numpy as np
import os
import pydicom as dicom
import pandas as pd
import shutil


# 获取dicom文件的pixel_array
def get_pixel_array(dicom_file_path):
    dcm_ori = dicom.dcmread(dicom_file_path)
    arr = dcm_ori.pixel_array
    arr = arr.astype(np.int32)
    return arr


# 最大最小归一化
def pixel_array_normalize(pixel_array):
    return 255.0 * (pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min())


def get_subtraction(dicom1, dicom2):
    """
    获取剪影图像
    :param dicom1:  增强后的dicom文件路径
    :param dicom2:  增强前的dicom文件路径
    :return:        剪影图像矩阵
    """
    pa1 = get_pixel_array(dicom1)
    pa2 = get_pixel_array(dicom2)
    pa1 = pixel_array_normalize(pa1)
    pa2 = pixel_array_normalize(pa2)
    img = pa1 - pa2

    # 截断负值
    img = np.clip(img, a_min=0, a_max=img.max())
    img = pixel_array_normalize(img)
    # 计算直方图
    img = np.uint8(img)
    hist = cv2.calcHist([img], [0], None, [255], [1, 256])
    # plt.plot(hist)
    # plt.show()
    hist_sum = np.cumsum(hist)
    # 截断归一化
    min_val = 0
    # min_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.01, side='left')
    max_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.9998, side='left')
    img = np.clip(img, a_min=min_val, a_max=max_val)
    img = 255.0 * (img - min_val) / (max_val - min_val)

    return img


# 获取duke数据集的剪影图像
def duke2subtraction(dir_path, save_path):
    caseId2sliceId = get_center_slice('E:\\MRI\\Duke\\Annotation_Boxes.xlsx')
    count = 0
    for case in os.listdir(dir_path):
        for tmp_dir in os.listdir(os.path.join(dir_path, case)):
            pre_sequence_dir = None
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                # if 'pre' in sequence_dir or 'ax 3d dyn' in sequence_dir or 'ax dynamic' in sequence_dir or 'ax dyn' in sequence_dir:
                if 'pre' in sequence_dir:
                # if '-ax 3d dyn' in sequence_dir:
                    pre_sequence_dir = sequence_dir
                    break
            if pre_sequence_dir is None:
                print('case: ' + case + ' pre sequence is not found.')
                continue
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                if '2nd' in sequence_dir:
                # if '2ax' in sequence_dir:
                    count += 1
                    sliceId = caseId2sliceId[case]
                    if sliceId < 100:
                        sliceId = '0' + str(sliceId)
                    else:
                        sliceId = str(sliceId)
                    dicom_file = '1-' + sliceId + '.dcm'
                    post_path = os.path.join(dir_path, case, tmp_dir, sequence_dir, dicom_file)
                    pre_path = os.path.join(dir_path, case, tmp_dir, pre_sequence_dir, dicom_file)
                    img = get_subtraction(post_path, pre_path)

                    cv2.imwrite(
                        os.path.join(save_path, case + '_second_subtraction_norm_' + dicom_file[2:5] + '.jpg'), img)
    print('序列数：', count)


# 转换duke数据集
def duke2jpg(dir_path, save_path):
    count = 0
    caseId2sliceId = get_center_slice('/data1/home/liukai/projects/TransFuseForBreast/dataprepare/Duke/dukeExample/Annotation_Boxes.xlsx')
    for case in os.listdir(dir_path):
        if not os.path.exists(os.path.join(save_path, case)):
            os.makedirs(os.path.join(save_path, case))
        for tmp_dir in os.listdir(os.path.join(dir_path, case)):
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                # if '2nd' in sequence_dir:
                if '2ax' in sequence_dir:
                    count += 1
                    sequence_save_path = os.path.join(save_path, case, sequence_dir)
                    if not os.path.exists(sequence_save_path):
                        os.makedirs(sequence_save_path)

                    # 仅获取中心切片图像
                    sliceId = caseId2sliceId[case]
                    if sliceId < 100:
                        sliceId = '0' + str(sliceId)
                    else:
                        sliceId = str(sliceId)
                    dicom_file = '1-' + sliceId + '.dcm'
                    pixel_array = get_pixel_array(os.path.join(dir_path, case, tmp_dir, sequence_dir, dicom_file))
                    pixel_array = pixel_array_normalize(pixel_array)
                    cv2.imwrite(os.path.join(save_path,
                                             case + '_second_post_' + dicom_file[2:5] + '.jpg'), pixel_array)

                    # for dicom_file in os.listdir(os.path.join(dir_path, case, tmp_dir, sequence_dir)):
                    #     pixel_array = get_pixel_array(os.path.join(dir_path, case, tmp_dir, sequence_dir, dicom_file))
                    #     # 归一化
                    #     pixel_array = pixel_array_normalize(pixel_array)
                    #     cv2.imwrite(os.path.join(sequence_save_path, dicom_file[2:5] + '.jpg'), pixel_array)
    print('序列数：', count)


# 根据是否含病灶进行分类
def duke_classify(annotation_path, dir_path, save_path):
    if not os.path.exists(os.path.join(save_path, 'positive')):
        os.makedirs(os.path.join(save_path, 'positive'))
    if not os.path.exists(os.path.join(save_path, 'negative')):
        os.makedirs(os.path.join(save_path, 'negative'))

    annotations = pd.read_excel(annotation_path)
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_slice = annotation[6]
        end_slice = annotation[7]
        case_path = os.path.join(dir_path, case_id)
        for sequence in os.listdir(case_path):
            sequence_path = os.path.join(case_path, sequence)
            for img in os.listdir(sequence_path):
                slice_num = int(img[0:3])
                if start_slice <= slice_num <= end_slice:
                    shutil.copyfile(os.path.join(sequence_path, img),
                                    os.path.join(save_path, 'positive', case_id + '_second_post_' + img))
                else:
                    shutil.copyfile(os.path.join(sequence_path, img),
                                    os.path.join(save_path, 'negative', case_id + '_second_post_' + img))


def get_center_slice(annotation_path):
    """
    获取每个病例病灶中心切片号
    :param annotation_path: 标注文件路径
    :return: 字典 key为病例ID value为切片号
    """
    annotations = pd.read_excel(annotation_path)
    caseId2sliceId = {}
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_slice = annotation[6]
        end_slice = annotation[7]
        caseId2sliceId[case_id] = int((end_slice - start_slice) / 2 + start_slice)
    return caseId2sliceId


def draw_box(img_dir, annotation_path):
    """
    画出标注框
    :param img_dir: 切片图片文件夹路径
    :param annotation_path: 标注文件路径
    :return: None
    """
    annotations = pd.read_excel(annotation_path)
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_row = annotation[2]
        start_column = annotation[4]
        end_row = annotation[3]
        end_column = annotation[5]
        start_slice = annotation[6]
        end_slice = annotation[7]
        sliceId = int((end_slice - start_slice) / 2 + start_slice)
        if sliceId < 100:
            sliceId = '0' + str(sliceId)
        else:
            sliceId = str(sliceId)
        img_name = os.path.join(img_dir, case_id + '_second_subtraction_norm3_' + sliceId + '.jpg')
        if os.path.exists(img_name):
            img = cv2.imread(img_name)
            # img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            cv2.rectangle(img, (start_column, start_row), (end_column, end_row), (0, 255, 0), 1)
            cv2.imwrite(os.path.join(img_dir, img_name[:-4] + '_draw.jpg'), img)


if __name__ == '__main__':
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # get_center_slice(annotation_path)

    # 画框
    # save_path = 'E:\\MRI\\Duke\\tmp'
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # draw_box(save_path, annotation_path)

    # 获取duke数据集的subtraction图像
    # dir_path = 'E:\\MRI\\Duke\\manifest-1680071275430\\Duke-Breast-Cancer-MRI'
    # save_path = 'E:\\MRI\\Duke\\subtractions'
    # duke2subtraction(dir_path, save_path)

    # 切片根据是否含病灶进行分类
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # save_path = 'E:\\MRI\\Duke'
    # dir_path = 'E:\\MRI\\Duke\\case'
    # duke_classify(annotation_path, dir_path, save_path)

    # 转换duke数据集
    dir_path = '/data1/home/liukai/projects/TransFuseForBreast/dataprepare/Duke/dukeExample/Duke-Breast-Cancer-MRI'
    save_path = '/data1/home/liukai/projects/TransFuseForBreast/dataprepare/Duke/dukeExample/second_post_center_norm'
    duke2jpg(dir_path, save_path)
